#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：domain-drop 
@File    ：custom_list_dataset.py
@IDE     ：PyCharm 
@Author  ：cao xu
@Date    ：2025/9/5 下午3:18 
"""
import os
from pathlib import Path
from PIL import Image, ImageFile, UnidentifiedImageError
ImageFile.LOAD_TRUNCATED_IMAGES = True    # 允许加载截断图像
Image.MAX_IMAGE_PIXELS = None             # 取消像素数限制（防止 DecompressionBombWarning）
from torch.utils.data import Dataset
IMAGE_SIZE = 224


def infer_hospital_from_path(img_path: str, images_root: str):
    # 约定：images_root/<医院>/<病例>/.../img.jpg
    p = Path(img_path)
    root = Path(images_root)
    try:
        rel = p.relative_to(root).as_posix()
        return rel.split("/")[0]
    except Exception:
        # 兜底：从路径中找 "images/<医院>/..."
        parts = p.as_posix().split("/")
        for i, s in enumerate(parts[:-1]):
            if s.lower() == "images" and i + 1 < len(parts):
                return parts[i + 1]
    return None


class ListTxtDataset(Dataset):
    """
    读取形如 "<相对images_root路径> <label_id>" 的清单；或清单里是绝对路径也可。
    仅保留 allowed_domains 内的样本；域ID由 dom2id 决定。
    """

    def __init__(self, list_file, images_root=None, transform=None,
                 allowed_domains=None, dom2id=None):
        self.root = images_root
        self.transform = transform
        self.allowed_domains = set(allowed_domains) if allowed_domains else None
        self.dom2id = dict(dom2id) if dom2id else {}
        self.samples = []

        with open(list_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                path_text, label_text = line.rsplit(" ", 1)
                label = int(label_text)

                # 绝对/相对路径均支持
                if not os.path.isabs(path_text) and self.root:
                    img_path = os.path.join(self.root, path_text)
                else:
                    img_path = path_text

                hosp = infer_hospital_from_path(img_path, self.root or "/")
                if self.allowed_domains and hosp not in self.allowed_domains:
                    continue
                if self.allowed_domains is None:
                    dom = self.dom2id.get(hosp, 0)  # 测试时 dom 随便给个 0 不影响
                else:
                    if hosp not in self.dom2id:
                        continue
                    dom = self.dom2id[hosp]
                self.samples.append((img_path, label, dom))

    def __len__(self):
        return len(self.samples)

    def _safe_open_rgb(self, path, max_retries=3):
        last_err = None
        for _ in range(max_retries):
            try:
                with Image.open(path) as im:
                    return im.convert("RGB")
            except (OSError, UnidentifiedImageError) as e:
                last_err = e
        # 记录坏样本到日志文件，便于后续清理
        badlog = os.path.join(os.path.dirname(__file__), "bad_images.log")
        try:
            with open(badlog, "a", encoding="utf-8") as f:
                f.write(f"{path}\t{repr(last_err)}\n")
        except Exception:
            pass
        # 彻底打不开：返回一个纯黑占位图，避免中断训练
        return Image.new("RGB", (IMAGE_SIZE if 'IMAGE_SIZE' in globals() else 224,
                                 IMAGE_SIZE if 'IMAGE_SIZE' in globals() else 224), color=0)

    def __getitem__(self, idx):
        img_path, label, dom = self.samples[idx]
        img = self._safe_open_rgb(img_path)
        if self.transform is not None:
            img = self.transform(img)
        return img, label, dom


class WithIndex(Dataset):
    """把 Dataset 包装为 ((img, y, d), idx) 以适配 train_domain.py 的取样格式。"""

    def __init__(self, base):
        self.base = base

    def __len__(self):
        return len(self.base)

    def __getitem__(self, i):
        x, y, d = self.base[i]
        return (x, y, d), i
